import os,pdb
try:
    from mpi4py import MPI
except ImportError:
    print 'cannot import MPI, continuing..'
from numpy import array,dot,zeros,shape,cross,zeros_like,sort,save,load
from numpy import round as npround
#from numpy.linalg import inv,cond,det,eigvals,eig,matrix_rank
from numpy.linalg import *
from scipy.linalg import lu
from math import sqrt
from v3dsolver.fluid_elements.fluid_elems import VortexPanel,FreeStream,Saffman
from v3dsolver.fluid_elements.fluid_elems import BoundaryBlob,ReleaseBlob
from v3dsolver.tree_manager.tree_manager import TreeManager
from v3dsolver.tree_manager.vel_calculator import VelCalculator
from v3dsolver.vector import Vec
from v3dsolver.tree_manager.distribution_manager import get_domain,DistributionManager
from v3dsolver.writer import write

class Geometry:
    """
    Generic class representing any general geometry. A geometry can have multi-
    -ple body objects in it.
    """
    def __init__(self,bodies):
        """
        Initialization for Geometry class. 
        
        bodies: list of body objects forming the geometry
        """
        self.bodies = bodies
        self.panels = self.get_panels() # contains panels of current proc only
        self.pts = self.get_pts() # contains pts of current proc only
    
    def get_panels(self):
        """
        Returns the panels contained in present Geometry object.
        """
        # total panels of all the bodies contained by Geometry object
        totalpanels = []
        for body in self.bodies:
            for panel in body.panels:
                totalpanels += [panel]
        return totalpanels
        
    def get_pts(self):
        """
        Returns the pts contained in present Geometry object.
        """
        # total pts of all the bodies contained by Geometry object
        totalpts = []
        for body in self.bodies:
            for pt in body.pts:
                totalpts += [pt]
        return totalpts

class Body:
    """
    Generic class representing a body. A body is a single 'connected' 
    entity which contains vortex panels inside it. 
    
    A body can be defined by a list of points which act as vertices of 
    vortex panels and are so ordered such that going around the vertices 
    of a panel would generate a normal which is 'outward' to the 
    specified body. The list of points can be inputted using a vtk file.
    """
    def __init__(self,filename,blob_rad,delta,vel=Vec()):
        """
        Constructor for Body object.
        
        filename: string of the form 'file.vtk' where file.vtk is stored
                  in data directory
        delta: distance of released blob from the boundary
        """
        self.filename = filename
        self.blob_rad = blob_rad
        self.vel = vel  # velocity of the body
        self.delta = delta
        self.pts = self.get_pts(filename)
        self.panels = self.get_panels(filename)
    
    def get_pts(self,filename):
        """
        Returns a list of points which make up the body object
        """
        filename = os.path.join(os.path.abspath('../../data'),filename)
        f = open(filename,'r+')
        lines = f.readlines()
        f.close()
        npts = int(lines[4].split()[1])
        pts = []
        for i,line in enumerate(lines[5:]):
            if len(pts)>=npts:
                break
            else:
                ptcoords = line.split()
                pt = []
                for j in range(len(ptcoords)):
                    pt += [float(ptcoords[j])]
                    if (j+1)%3==0:
                        pts += [array(pt)]
                        pt = []
        return pts
    
    def get_panels(self,filename):
        """
        Returns a list of VortexPanel objects possesed by this body object, 
        based on points and panels data in file named 'filename'.
        
        delta: distance of to-be released blobs from panel surface
        """
        filename = os.path.join(os.path.abspath('../../data'),filename)
        f = open(filename,'r+')
        lines = f.readlines()
        f.close()
        currentline = len(self.pts)/3 + 5
        panels = []
        npanels = int(lines[currentline+1].split()[1])
        for line in lines[currentline+2:]:
            if len(panels)==npanels:
                break
            indices = line.split()[1:]
            vertices = []
            for index in indices:
                index = int(index)
                vertices += [self.pts[index]]
            panels += [VortexPanel(vertices,self.blob_rad,self.vel,
                                   self.delta)]
        return panels

class BoundaryManager:
    """
    Class representing the boundary manager object which deals with definition
    of the geometry, satisfaction of no-penetration and no-slip boundary 
    conditions and finally releasing the blobs to the fluid domain.
    """
    def __init__(self,geometry,blob_rad,freestream,this_outsideblobs=[]):
        """
        Constructor of BoundaryManager object. Defines the geometry and panels
        associated with them based on point and panel data obtained from the 
        vtk file named 'filename' stored in 'data'.
        
        geometry: Geometry object consisting of a single body/list of bodies
        blob_rad: radius of each blob
        freestream: object of type fluid_elems.FreeStream specifying existing
                    freestream conditions
        this_outsideblobs: blobs already existing in flow domain with current
                           processor
        """
        self.comm = MPI.COMM_WORLD
        self.blob_rad = blob_rad
        self.geometry = geometry
        self.globalpanels = self.geometry.panels
        self.pts = self.geometry.pts
        dm = DistributionManager(self.comm,self.globalpanels)
        self.panels = dm.get_this_proc_elems()
        
        # get back global panels
        global_panels = self.comm.allgather(self.panels)
        self.globalpanels = []        
        for i in range(len(global_panels)):
            self.globalpanels += global_panels[i]
        
        
        self.np_controlpts = self.get_npcontrolpts(self.panels)
        self.ns_controlpts = self.get_nscontrolpts(self.panels)
        
        #self.coeffs1 = self.get_coeffs1()
        #print 'generating pinv:'
        #self.pinv_coeffs = pinv(self.coeffs1)
        #save('pinv_coeff',self.pinv_coeffs)
        self.pinv_coeffs = load('pinv_coeff.npy')
        self.freestream = freestream
        self.this_outsideblobs = this_outsideblobs
        self.this_bblobs = []
    
    def get_npcontrolpts(self,panels):
        """
        Returns a list of controlpts contained by panels on which no 
        penetration will be satisfied. 
        """
        npcontrolpts = []
        for panel in panels:
            npcontrolpts += panel.np_controlpts
        return npcontrolpts
    
    def get_nscontrolpts(self,panels):
        """
        Returns a list of controlpts contained by panels on which no 
        slip will be satisfied. 
        """
        nscontrolpts = []
        for panel in panels:
            nscontrolpts += [panel.ns_controlpt]
        return nscontrolpts

    def get_coeffs(self):
        """
        Returns coefficient matrix and it's inverse containing velocity 
        contributions due to kernel cause_elems on controlpts.
        """
        coeffs = zeros((4*len(self.panels),4*len(self.panels)))
        print 'generating matrix..'
        for i,cause_panel in enumerate(self.panels):
            for j,cause_blob in enumerate(cause_panel.cause_elems):
                for k in range(2):
                    vordirn = cause_panel.basis[k]
                    kernel_cause_elem = Saffman(pos=cause_blob.pos,vor=vordirn,
                                          vel=cause_blob.vel,rad=cause_blob.rad)
                    
                    for l,effect_panel in enumerate(self.panels):
                        # 2 equations of no penetration at 2 control pts
                        for m in range(2):
                            effect_elem = effect_panel.np_controlpts[m]
                            coeffs[l*4+m][i*4+j*2+k] = kernel_cause_elem.get_vel(effect_elem).dot(effect_panel.unitnormal)
                        # 2 equations of no slip at 1 control pt
                        effect_elem = effect_panel.ns_controlpt
                        for n in range(2):
                            coeffs[l*4+2+n][i*4+j*2+k] = kernel_cause_elem.get_vel(effect_elem).dot(effect_panel.basis[n])
                    # making sum of vorticity in x,y,z directions zero
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(1,0,0))
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(0,1,0))
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(0,0,1))
        save('A_sphere_Koumoutsakos',coeffs)
        print 'matrix generated'
        inv_coeffs = inv(coeffs)
        print 'matrix inverted'
        save('inv_A_sphere_Koumoutsakos',inv_coeffs)
        return coeffs,inv_coeffs
        
    def get_coeffs1(self):
        """
        Returns coefficient matrix and it's inverse containing velocity 
        contributions due to kernel cause_elems on controlpts.
        """
        coeffs = zeros((5*len(self.globalpanels),4*len(self.globalpanels)))
        print 'generating matrix1..'
        for i,cause_panel in enumerate(self.globalpanels):
            for j,cause_blob in enumerate(cause_panel.cause_elems):
                for k in range(2):
                    vordirn = cause_panel.basis[k]
                    kernel_cause_elem = Saffman(pos=cause_blob.pos,vor=vordirn,
                                          vel=cause_blob.vel,rad=cause_blob.rad)
                    
                    for l,effect_panel in enumerate(self.globalpanels):
                        # 3 equations of no penetration at 3 control pts
                        for m in range(3):
                            effect_elem = effect_panel.np_controlpts[m]
                            coeffs[l*5+m][i*4+j*2+k] = kernel_cause_elem.get_vel(effect_elem).dot(effect_panel.unitnormal)
                        # 2 equations of no slip at 1 control pt
                        effect_elem = effect_panel.ns_controlpt
                        for n in range(2):
                            coeffs[l*5+3+n][i*4+j*2+k] = kernel_cause_elem.get_vel(effect_elem).dot(effect_panel.basis[n])
                    # making sum of vorticity in x,y,z directions zero
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(1,0,0))
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(0,1,0))
                    #coeffs[2*len(self.panels)][i*4+j*2+k] = cause_panel.basis[k].dot(Vec(0,0,1))
        save('A_sphere1_Koumoutsakos',coeffs)
        print 'matrix1 generated'
        #inv_coeffs = inv(coeffs)
        #print 'matrix inverted'
        #save('inv_A_sphere_Koumoutsakos',inv_coeffs)
        return coeffs#,inv_coeffs
    
    def rhs(self):
        """
        Returns rhs for solving boundary conditions at controlpts. Contains
        contributions from freestream and outsideblobs. Works only for 
        serial version for now.
        """
        # computing velocity on control points
        elems = self.np_controlpts + self.ns_controlpts + self.this_outsideblobs
        domain = get_domain(elems,findglobal=True) # global domain calculation
        tree = TreeManager(elems,domain,npmax=5,npseudo=16,rad=self.blob_rad,
                           comm=self.comm)
        tree.root.effect_elems = self.ns_controlpts + self.np_controlpts #ignoring outside effect elems
        VelCalculator(tree).parallel_vel_computation()
        for controlpt in self.np_controlpts+self.ns_controlpts:
            controlpt.vel += self.freestream.get_vel(controlpt)
        
        # building rhs
        rhs_vec = zeros(4*len(self.panels))
        for l,effect_panel in enumerate(self.panels):
            for m in range(2):
                effect_elem = effect_panel.np_controlpts[m]
                print 'effect_vel:', effect_elem.vel
                rhs_vec[l*4+m] = (effect_panel.vel - effect_elem.vel).dot(effect_panel.unitnormal)
            effect_elem = effect_panel.ns_controlpt
            for n in range(2):
                rhs_vec[l*4+2+n] = (effect_panel.vel - effect_elem.vel).dot(effect_panel.basis[n])
        #rhs_vec[2*len(self.panels)] = 0.
        #rhs_vec[2*len(self.panels)] = 0.
        #rhs_vec[2*len(self.panels)] = 0.
        print 'rhs_vec:', rhs_vec
        return rhs_vec
    
    def rhs1(self):
        """
        Returns rhs for solving boundary conditions at controlpts. Contains
        contributions from freestream and outsideblobs. Works only for 
        serial version for now.
        """
        # computing velocity on control points
        elems = self.np_controlpts + self.ns_controlpts + self.this_outsideblobs
        domain = get_domain(elems,findglobal=True) # global domain calculation
        tree = TreeManager(elems,domain,npmax=5,npseudo=16,rad=self.blob_rad,
                           comm=self.comm)
        tree.root.effect_elems = self.ns_controlpts + self.np_controlpts #ignoring outside effect elems
        VelCalculator(tree).parallel_vel_computation()
        for controlpt in self.np_controlpts+self.ns_controlpts:
            controlpt.vel += self.freestream.get_vel(controlpt)
        
        # building local rhs
        local_rhs = zeros(5*len(self.panels))
        for l,effect_panel in enumerate(self.panels):
            for m in range(3):
                effect_elem = effect_panel.np_controlpts[m]
                local_rhs[l*5+m] = (effect_panel.vel - effect_elem.vel).dot(effect_panel.unitnormal)
            effect_elem = effect_panel.ns_controlpt
            for n in range(2):
                local_rhs[l*5+3+n] = (effect_panel.vel - effect_elem.vel).dot(effect_panel.basis[n])
        local_rhs = list(local_rhs)
        
        # combine all local rhs
        global_rhs = self.comm.allgather(local_rhs)
        global_rhs_new = []        
        for i in range(len(global_rhs)):
            global_rhs_new += global_rhs[i]

        global_rhs_new = array(global_rhs_new)
        return global_rhs_new
    
    def solve_bc(self):
        """
        Solves no slip and no penetration boundary conditions by equating 
        velocity at one controlpt per panel to panel velocity.
        """
        b = self.rhs()
        vors = dot(self.inv_coeffs,b)
        #vors = lstsq(self.coeffs,b)[0]
        #vors = solve(self.coeffs,b)
        #print 'vors:', vors
        #print 'b:', b
        #print 'coeffs:', coeffs
        for i,cause_panel in enumerate(self.panels):
            for j,cause_blob in enumerate(cause_panel.cause_elems):
                vor = Vec()
                for k in range(2):
                    vor += vors[i*4+j*2+k]*cause_panel.basis[k]
                cause_blob.vor = vor
        write(self.np_controlpts,'np_controlpts')
        
        elems = []
        for panel in self.panels:
            for vertex in panel.vertices:
                elem = Saffman(pos=vertex,vel=Vec(),vor=Vec(),rad=panel.rad)
                elems += [elem]
        write(elems,'vertices')
        
    def solve_bc1(self):
        """
        Solves no slip and no penetration boundary conditions by equating 
        velocity at one controlpt per panel to panel velocity.
        """
        b = self.rhs1()
        vors = dot(self.pinv_coeffs,b)
        
        for i,cause_panel in enumerate(self.globalpanels):
            for j,cause_blob in enumerate(cause_panel.cause_elems):
                vor = Vec()
                for k in range(2):
                    vor += vors[i*4+j*2+k]*cause_panel.basis[k]
                cause_blob.vor = vor
        #write(self.np_controlpts,'np_controlpts')
        #write(self.ns_controlpts,'ns_controlpts')
    
    def verify_bc(self,bblobs,rblobs):
        """
        Computes velocities on controlpts due to obtained vorticity 
        distribution and checks if the boundary conditions are 
        accurately satisfied.
        """
        # compute all velocities
        elems = bblobs + rblobs + self.np_controlpts + self.ns_controlpts + self.this_outsideblobs
        VelCalculator().direct_vel_computation(elems)
        for elem in elems:
            elem.vel += self.freestream.get_vel(elem)
        
        # print velocities
        for i, panel in enumerate(self.panels):
            for np_controlpt in panel.np_controlpts:
                print i, 'np:', np_controlpt.vel.dot(panel.unitnormal)
            for k in range(2):
                print i, 'ns:', panel.ns_controlpt.vel.dot(panel.basis[k])
    
    def get_released_particles(self,this_outsideblobs):
        """
        Returns two lists: one for boundary elems; another for released elems 
        obtained after satisfying no slip and no penetration conditions.
        
        this_outsideblobs: list of blobs of current processor outside the 
                           boundary
        """
        # solve boundary conditions
        self.this_outsideblobs = this_outsideblobs
        self.solve_bc1()
        
        # collect boundary blobs and released blobs
        bblobs,rblobs,outsideblobs = [],[],[]
        for panel in self.globalpanels:
            bblobs += [panel.bblob]
            rblobs += [panel.rblob]
            outsideblobs += [Saffman(pos=panel.rblob.pos,vor=panel.rblob.vor,
                                     vel=panel.rblob.vel,rad=panel.rblob.rad)]
        
        this_bblobs = DistributionManager(self.comm,bblobs).get_this_proc_elems()
        this_rblobs = DistributionManager(self.comm,rblobs).get_this_proc_elems()
        this_new_outsideblobs = DistributionManager(self.comm,outsideblobs).get_this_proc_elems()
        self.this_bblobs = this_bblobs 
        
        # verify obtained velocities on controlpts
        #self.verify_bc(bblobs,rblobs)
        #write(rblobs,'sphere_rblobs_Koumoutsakos')
        #write(bblobs,'sphere_bblobs_Koumoutsakos')
        #pdb.set_trace()
        
        return this_bblobs,this_rblobs,this_new_outsideblobs

if __name__=='__main__':
    b = Body('sphere.vtk',0.005)
    panels = b.panels
    print len(panels)
